from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
import os

class RAGSystem:
    VECTOR_DB_DIR = "./chroma_db"
    
    @classmethod
    def get_vector_store(cls):
        """获取或创建向量存储"""
        embeddings = cls._get_embeddings()
        
        # 检查是否已有持久化的向量库
        if cls._vector_db_exists():
            print("检测到已存在的向量库，直接加载...")
            return Chroma(
                persist_directory=cls.VECTOR_DB_DIR,
                embedding_function=embeddings
            )
        
        # 创建新的向量库
        return cls._create_new_vector_store(embeddings)
    
    @staticmethod
    def _get_embeddings():
        """创建嵌入模型"""
        return HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={'device': 'cpu'}
        )
    
    @classmethod
    def _vector_db_exists(cls):
        """检查向量库是否已存在"""
        return os.path.exists(cls.VECTOR_DB_DIR) and os.listdir(cls.VECTOR_DB_DIR)
    
    @classmethod
    def _create_new_vector_store(cls, embeddings):
        """创建新的向量存储"""
        # 检查并加载文档
        documents = cls._load_documents()
        
        # 分割文档
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        splits = text_splitter.split_documents(documents)
        
        # 创建并持久化向量存储
        return Chroma.from_documents(
            documents=splits,
            embedding=embeddings,
            persist_directory=cls.VECTOR_DB_DIR
        )
    
    @classmethod
    def _load_documents(cls):
        """加载所有文档"""
        files_dir = "files"
        if not os.path.exists(files_dir):
            os.makedirs(files_dir, exist_ok=True)
            return [Document(
                page_content="默认文档内容",
                metadata={"source": "default"}
            )]
        
        loaders = []
        
        # 加载文本文件
        if any(f.endswith('.txt') for f in os.listdir(files_dir)):
            loaders.append(DirectoryLoader(
                files_dir,
                glob="**/*.txt",
                loader_cls=TextLoader,
                loader_kwargs={"encoding": "utf-8"}
            ))
        
        # 加载PDF文件
        pdf_files = [f for f in os.listdir(files_dir) if f.endswith('.pdf')]
        for pdf_file in pdf_files:
            pdf_path = os.path.join(files_dir, pdf_file)
            try:
                loaders.append(PyPDFLoader(pdf_path))
            except Exception as e:
                print(f"加载PDF文件 {pdf_file} 时出错: {e}")
        
        # 合并所有文档
        all_docs = []
        for loader in loaders:
            try:
                docs = loader.load()
                all_docs.extend(docs)
            except Exception as e:
                print(f"加载文档时出错: {e}")
        
        return all_docs if all_docs else [Document(
            page_content="默认文档内容",
            metadata={"source": "default"}
        )]
